// This program demonstrates File Accessed(open) By Container(aka. fabc).

package main

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"os"
	"os/signal"
	"path"
	"sort"
	"sync"
	"syscall"
	"time"
	"unsafe"

	"github.com/cilium/ebpf/link"
	"github.com/cilium/ebpf/ringbuf"
	"github.com/cilium/ebpf/rlimit"
	"golang.org/x/sys/unix"
)

//go:generate go run github.com/cilium/ebpf/cmd/bpf2go -target $GOARCH -type event bpf fabc.c -- -I$ebpf_headers
const (
	SizeKB = 1024
	SizeMB = 1024 * 1024
	SizeGB = 1024 * 1024 * 1024
)

type Logger struct {
	sync.Mutex
	logs  map[string]*os.File
	count uint64
}

func NewLogger() *Logger {
	return &Logger{
		logs: make(map[string]*os.File),
	}
}

func (log *Logger) Printf(fmts string, args ...interface{}) {
	fmt.Printf(fmts+"\n", args...)
}

func (log *Logger) Fatalf(fmts string, args ...interface{}) {
	fmt.Printf(fmts+"\n", args...)
	os.Exit(-1)
}

func (log *Logger) Log(cid string, fmts string, args ...interface{}) {
	log.Lock()
	defer log.Unlock()
	os.MkdirAll("logs", os.ModePerm)
	f, ok := log.logs[cid]
	if !ok {
		f, err := os.OpenFile("logs/"+cid+".log", os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0600)
		if err != nil {
			return
		}
		log.logs[cid] = f
	}
	f.WriteString(fmt.Sprintf(fmts+"\n", args...))
	log.count += 1
	if log.count%1000 == 0 {
		go log.Flush()
	}
}

func (log *Logger) Flush() {
	log.Lock()
	defer log.Unlock()
	for _, f := range log.logs {
		f.Sync()
	}
}

var logFirstOnly = true // 只打印首次访问的文件，后续访问不打印
var logPseudoFS = false // 打印伪文件系统中的文件
var log = NewLogger()

type FileInfo struct {
	*bpfEvent
	cid   string
	comm  string
	path  string
	count uint64
}

func NewFileInfo(e *bpfEvent) *FileInfo {
	fi := &FileInfo{
		bpfEvent: e,
		cid:      unix.ByteSliceToString(e.Cid[:]),
		comm:     unix.ByteSliceToString(e.Comm[:]),
		path:     unix.ByteSliceToString(e.Path[:]),
	}
	if len(fi.cid) > 10 {
		fi.cid = fi.cid[:10]
	}
	if len(fi.comm) > 20 {
		fi.comm = fi.comm[:20]
	}
	return fi
}

func (fi *FileInfo) Key() uint64 {
	return uint64(10000000*fi.Pid) + fi.Timestamp
}

type ContainerInfo struct {
	ts      uint64
	cid     string
	parents map[uint64]*FileInfo
	files   map[string]*FileInfo
}

func NewContainerInfo(fi *FileInfo) *ContainerInfo {
	ci := ContainerInfo{
		ts:      fi.Timestamp,
		cid:     fi.cid,
		parents: make(map[uint64]*FileInfo),
		files:   make(map[string]*FileInfo),
	}
	ci.Add(fi)
	return &ci
}

func (ci *ContainerInfo) Add(fi *FileInfo) {
	if fi.EventType == 1 {
		p, ok := ci.parents[fi.Key()]
		if ok {
			p.path = path.Join(fi.path, p.path)
		} else {
			ci.parents[fi.Key()] = fi
		}

		return
	}

	p, ok := ci.parents[fi.Key()]
	if ok {
		fi.path = path.Join(p.path, fi.path)
	}
	delete(ci.parents, fi.Key())

	if old, ok := ci.files[fi.path]; ok {
		old.count += 1
		ci.Log(old)
	} else {
		fi.count += 1
		ci.files[fi.path] = fi
		ci.Log(fi)
	}
}

func (ci *ContainerInfo) Log(fi *FileInfo) {
	// https://github.com/torvalds/linux/blob/master/include/uapi/linux/magic.h
	mfs := map[uint64]string{
		0x62656572: "sysfs",
		0x9fa0:     "procfs",
		0x01021994: "tmpfs",
		0x1cd1:     "devpts",
		0x27e0eb:   "cgroupfs",
		0x794c7630: "overlayfs",
		0xef53:     "extfs",
	}
	fs, ok := mfs[fi.Magic]
	if !ok {
		fs = fmt.Sprintf("0x%x", fi.Magic)
	}
	size := ""
	if fi.Size < SizeKB {
		size = fmt.Sprintf("%4dB", fi.Size)
	} else if fi.Size < SizeMB {
		size = fmt.Sprintf("%3dKB", fi.Size/SizeKB)
	} else if fi.Size < SizeGB {
		size = fmt.Sprintf("%3dMB", fi.Size/SizeMB)
	} else {
		size = fmt.Sprintf("%3dGB", fi.Size/SizeMB)
	}

	if fi.count > 1 && logFirstOnly {
		return
	}

	if !logPseudoFS && (fs == "sysfs" || fs == "procfs" || fs == "cgroupfs") {
		return
	}

	log.Log(fi.cid, "time=%v cid=%v pid=%-8d comm=%-20s : %s x %v %v:%v",
		time.Now().Format("2006-01-02 15:04:05.000"), ci.cid, int(fi.Pid), fi.comm,
		size, fi.count, fs, fi.path,
	)
}

type AccessInfo struct {
	sync.Mutex
	nc  map[uint64]string         // ns_id:cid
	cci map[string]*ContainerInfo // cid:ContainerInfo
}

func NewAccessInfo() AccessInfo {
	return AccessInfo{
		nc:  make(map[uint64]string),
		cci: make(map[string]*ContainerInfo),
	}
}

func (ai *AccessInfo) Add(e *bpfEvent) {
	ai.Lock()
	defer ai.Unlock()
	fi := NewFileInfo(e)

	if len(fi.cid) != 0 {
		ai.nc[fi.Mntns] = fi.cid
	} else { // process like nvidia-docker nsenter container to access some files.
		fi.cid = ai.nc[fi.Mntns]
	}
	if len(fi.cid) == 0 {
		return
	}

	ci, ok := ai.cci[fi.cid]
	if !ok {
		ci := NewContainerInfo(fi)
		if ci == nil {
			log.Printf("invalid event %#+v", fi)
			return
		}
		ai.cci[ci.cid] = ci
	} else {
		ci.Add(fi)
	}
}

func (ai *AccessInfo) PrintSummary() {
	ai.Lock()
	defer ai.Unlock()
	log.Printf("\033[H\033[2J")
	log.Printf("fabc: files accessed by container")
	log.Printf("container       total |  overlayfs      extfs")
	var keys sort.StringSlice = make([]string, 0, len(ai.cci))
	for k := range ai.cci {
		keys = append(keys, k)
	}
	keys.Sort()

	for _, k := range keys {
		ci := ai.cci[k]
		size := uint64(0)
		overlaySize := uint64(0)
		extSize := uint64(0)
		for _, fi := range ci.files {
			size += fi.Size
			if fi.Magic == 0x794c7630 {
				overlaySize += fi.Size
			} else if fi.Magic == 0xef53 {
				extSize += fi.Size
			}

		}
		log.Printf("%s %8dMB | %8dMB %8dMB",
			ci.cid, size/SizeMB, overlaySize/SizeMB, extSize/SizeMB)
		log.Log(ci.cid, "time=%v cid=%s total=%dMB overlayfs=%dMB extfs=%dMB",
			time.Now().Format("2006-01-02 15:04:05.000"), ci.cid, size/SizeMB, overlaySize/SizeMB, extSize/SizeMB)
	}
}

func main() {
	stopper := make(chan os.Signal, 1)
	signal.Notify(stopper, os.Interrupt, syscall.SIGTERM)

	// Allow the current process to lock memory for eBPF resources.
	if err := rlimit.RemoveMemlock(); err != nil {
		log.Fatalf("removing memlock: %v", err)
	}

	// Load pre-compiled programs and maps into the kernel.
	objs := bpfObjects{}
	if err := loadBpfObjects(&objs, nil); err != nil {
		log.Fatalf("loading objects: %v", err)
	}
	defer objs.Close()

	kp, err := link.Kprobe("copy_process", objs.ReportNs, nil)
	if err != nil {
		log.Fatalf("opening kprobe: %s", err)
	}
	defer kp.Close()

	kr, err := link.Kretprobe("do_filp_open", objs.ReportOpen, nil)
	if err != nil {
		log.Fatalf("opening kprobe: %s", err)
	}
	defer kr.Close()

	rd, err := ringbuf.NewReader(objs.bpfMaps.Events)
	if err != nil {
		log.Fatalf("opening ringbuf reader: %s", err)
	}
	defer rd.Close()
	ai := NewAccessInfo()

	go func() {
		ticker := time.NewTicker(10*time.Second)
		for _ = range ticker.C {
			ai.PrintSummary()
		}
	}()

	go func() {
		for {
			<-stopper
			log.Flush()
			if err := rd.Close(); err != nil {
				log.Fatalf("closing ringbuf reader: %s", err)
			}
		}
	}()

	for {
		record, err := rd.Read()
		if err != nil {
			if errors.Is(err, ringbuf.ErrClosed) {
				log.Printf("received signal, exiting..")
				return
			}
			log.Printf("reading from reader: %s", err)
			continue
		}

		var event bpfEvent
		var endian binary.ByteOrder
		if (*(*[2]uint8)(unsafe.Pointer(&[]uint16{1}[0])))[0] == 0 {
			endian = binary.BigEndian
		} else {
			endian = binary.LittleEndian
		}
		if err := binary.Read(bytes.NewBuffer(record.RawSample), endian, &event); err != nil {
			log.Printf("parsing ringbuf event: %s", err)
			continue
		}
		ai.Add(&event)
	}
}
